import torch
import matplotlib.pyplot as plt
import numpy as np

# 波束图计算函数
def compute_beam_pattern(af, theta, N, d, f, c=3e8):
    k = 2 * np.pi * f / c
    beam_pattern = 0
    for n in range(N):
        phase_shift = np.exp(1j * k * d * n * np.sin(theta))
        beam_pattern += af[n] * phase_shift
    return np.abs(beam_pattern)

# 绘制波束图
def plot_beam_pattern(af, theta_scan, N, d, f):
    theta_range = np.linspace(-90, 90, 360) * np.pi / 180
    beam_pattern = np.zeros_like(theta_range)
    
    for i, theta in enumerate(theta_range):
        beam_pattern[i] = compute_beam_pattern(af, theta, N, d, f)
    
    plt.figure(figsize=(10, 6))
    plt.plot(theta_range * 180 / np.pi, 20 * np.log10(beam_pattern / np.max(beam_pattern)))
    plt.xlabel('Theta (degrees)')
    plt.ylabel('Beam Pattern (dB)')
    plt.title('Beam Pattern of the Antenna Array')
    plt.grid(True)
    
    # 标记扫描角度
    plt.axvline(x=theta_scan, color='r', linestyle='--')
    plt.axhline(y=20 * np.log10(compute_beam_pattern(af, np.radians(theta_scan), N, d, f) / np.max(beam_pattern)), color='r', linestyle='--')
    plt.show()

# 参数设置
N = 32
d = 0.25
f = 300.0e6
theta_scan = 5  # 扫描角度

# 假设的天线因子AF，这里使用均匀分布的幅度和相位
af = np.ones(N)  # 假设所有天线幅度相等
phases = np.exp(1j * np.linspace(0, 2 * np.pi * (N - 1) / N, N))  # 均匀分布的相位
af = af * phases  # 结合幅度和相位

# 绘制波束图
plot_beam_pattern(af, theta_scan, N, d, f)
